import cv2
import rclpy
from cv_bridge import CvBridge
from rclpy.node import Node
from zorbus_interface.msg import RGBDImage


# ========== 图像订阅节点 ==========
class ImageDepthSubscriber(Node):
    def __init__(self):
        super().__init__('rgbd_subscriber')
        self.subscription = self.create_subscription(RGBDImage, 'rgbd_topic', self.listener_callback, 10)
        self.bridge = CvBridge()

    def listener_callback(self, msg):
        rgb = self.bridge.imgmsg_to_cv2(msg.rgb_image, desired_encoding='bgr8')
        depth = self.bridge.imgmsg_to_cv2(msg.depth_image, desired_encoding='16UC1')
        self.get_logger().info('Received RGBD image')
        cv2.imwrite('RGB.jpg', rgb)
        cv2.imwrite('Depth.jpg', depth.astype('float32') / 10000)

# ========== 主函数 ==========
def main(args=None):
    rclpy.init()
    rgbd_sub = ImageDepthSubscriber()
    rclpy.spin(rgbd_sub)
    rgbd_sub.destroy_node()
    rclpy.shutdown()
